import cv2
import numpy as np
import mediapipe as mp

# Mediapipe FaceMesh
mp_face_mesh = mp.solutions.face_mesh
cap = cv2.VideoCapture(0)

# 目标人脸
target_img = cv2.imread("face.png")
target_img = cv2.resize(target_img, (256, 256))

# 定义目标脸的三个关键点（大概位置）
target_pts = np.array([
    (70, 120),   # 左眼
    (190, 120),  # 右眼
    (130, 200)   # 嘴巴
], dtype=np.float32)

def color_transfer(src, dst):
    """把 src 的颜色匹配到 dst"""
    src_mean, src_std = cv2.meanStdDev(src)
    dst_mean, dst_std = cv2.meanStdDev(dst)

    # 转换成 (1,1,3)，方便广播
    src_mean, src_std = src_mean.reshape((1,1,3)), src_std.reshape((1,1,3))
    dst_mean, dst_std = dst_mean.reshape((1,1,3)), dst_std.reshape((1,1,3))

    dst = dst.astype(np.float32)
    dst = (dst - dst_mean) * (src_std / (dst_std + 1e-6)) + src_mean
    dst = np.clip(dst, 0, 255).astype(np.uint8)
    return dst


with mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, refine_landmarks=True) as face_mesh:
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        h, w, _ = frame.shape
        results = face_mesh.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

        if results.multi_face_landmarks:
            for face_landmarks in results.multi_face_landmarks:
                left_eye = face_landmarks.landmark[33]   # 左眼
                right_eye = face_landmarks.landmark[263] # 右眼
                mouth = face_landmarks.landmark[13]      # 嘴巴

                pts = np.array([
                    (int(left_eye.x * w), int(left_eye.y * h)),
                    (int(right_eye.x * w), int(right_eye.y * h)),
                    (int(mouth.x * w), int(mouth.y * h))
                ], dtype=np.float32)

                # 仿射变换
                M = cv2.getAffineTransform(target_pts, pts)
                warped_face = cv2.warpAffine(target_img, M, (w, h))

                # 提取人脸区域（mask）
                gray = cv2.cvtColor(warped_face, cv2.COLOR_BGR2GRAY)
                _, mask = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)

                # 模糊边缘，避免生硬
                mask = cv2.GaussianBlur(mask, (25, 25), 10)

                # 颜色匹配（目标脸 → 当前帧）
                face_region = cv2.bitwise_and(frame, frame, mask=mask)
                warped_face = color_transfer(face_region, warped_face)

                # 无缝融合
                center = (w // 2, h // 2)
                output = cv2.seamlessClone(warped_face, frame, mask, center, cv2.NORMAL_CLONE)

                cv2.imshow("Swap Result", output)
        else:
            cv2.imshow("Swap Result", frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

cap.release()
cv2.destroyAllWindows()
